#include "tos_secure_context.h"
#include "tos_secure_port_macros.h"
#include "tos_k.h"

#define K_SECURE_CTX_STACK_ALIGNMENT               8
#define K_SECURE_CTX_STACK_SEAL_SIZE               8
#define K_SECURE_CTX_STACK_SEAL_VALUE              0xFEF5EDA5

extern void tos_secure_load_ctx_asm( k_secure_ctx_t * task_secure_context ) __attribute__( ( naked ) );
extern void tos_secure_save_ctx_asm( k_secure_ctx_t * task_secure_context ) __attribute__( ( naked ) );

/**
 * @brief Maximum number of secure contexts.
 */
#ifndef TOS_CFG_NUM_SECURE_CTX
    #define TOS_CFG_NUM_SECURE_CTX    8UL
#endif

/**
 * @brief CONTROL value for privileged tasks.
 *
 * Bit[0] - 0 --> Thread mode is privileged.
 * Bit[1] - 1 --> Thread mode uses PSP.
 */
#define SECURE_CONTROL_VALUE_PRIVILEGED      0x02

/**
 * @brief CONTROL value for un-privileged tasks.
 *
 * Bit[0] - 1 --> Thread mode is un-privileged.
 * Bit[1] - 1 --> Thread mode uses PSP.
 */
#define SECURE_CONTROL_VALUE_UNPRIVILEGED    0x03

typedef enum secure_ctx_state {
    SECURE_CTX_UNINITIALIZED,
    SECURE_CTX_INITIALIZED
} secure_ctx_state_t;

static k_secure_ctx_t k_secure_ctxs[TOS_CFG_NUM_SECURE_CTX];

__STATIC__ k_secure_ctx_handle_t tos_secure_ctx_get(void *sec_task_handler)
{
    uint32_t i;
    k_secure_ctx_handle_t index = TOS_CFG_NUM_SECURE_CTX;
    for (i = 0; i < TOS_CFG_NUM_SECURE_CTX; i++) {
        if ((k_secure_ctxs[i].task_handler== K_NULL) && (index == TOS_CFG_NUM_SECURE_CTX)) {
            index = i;
        } else if (k_secure_ctxs[i].task_handler == sec_task_handler) {
            index = TOS_CFG_NUM_SECURE_CTX;
            break;
        }
    }

    return index;
}

__NSC__ __KNL__ void tos_secure_ctx_init(void)
{
    uint32_t i;
    uint32_t ulIPSR;
    static secure_ctx_state_t state = SECURE_CTX_UNINITIALIZED;

    secure_get_IPSR(ulIPSR);

    if (state == SECURE_CTX_UNINITIALIZED && ulIPSR != 0) {
        secure_set_PSP(K_NULL);
        secure_set_PSPLIM(K_NULL);
        for (i = 0; i < TOS_CFG_NUM_SECURE_CTX; i++) {
            k_secure_ctxs[i].curr_psp = K_NULL;
            k_secure_ctxs[i].psp_limit = K_NULL;
            k_secure_ctxs[i].psp_start = K_NULL;
            k_secure_ctxs[i].task_handler = K_NULL;
        }
        #if ( __MPU_PRESENT == 1 )
        {
            /* Configure thread mode to use PSP and to be unprivileged. */
            secure_set_CONTROL( SECURE_CONTROL_VALUE_UNPRIVILEGED );
        }
        #else /* configENABLE_MPU */
        {
            /* Configure thread mode to use PSP and to be privileged. */
            secure_set_CONTROL(SECURE_CONTROL_VALUE_PRIVILEGED);
        }
        #endif /* configENABLE_MPU */
        // mmheap_init();
        state = SECURE_CTX_INITIALIZED;
    }
}

__NSC__ __KNL__ k_secure_ctx_handle_t tos_secure_ctx_alloc(uint32_t size, void *task_handler)
{
    k_stack_t *stk_base = K_NULL;
    uint32_t ulIPSR;
    k_secure_ctx_handle_t index = K_SECURE_CTX_INVALID_ID;

    secure_get_IPSR(ulIPSR);
    if (ulIPSR != 0) {
        index = tos_secure_ctx_get(task_handler);
        if (index < TOS_CFG_NUM_SECURE_CTX) {
            stk_base = tos_mmheap_aligned_alloc(size + K_SECURE_CTX_STACK_SEAL_SIZE, K_SECURE_CTX_STACK_ALIGNMENT);
            if (stk_base != K_NULL) {
                k_secure_ctxs[index].curr_psp = (k_task_t*)(stk_base + size);
                k_secure_ctxs[index].psp_limit = (k_task_t*)stk_base;
                k_secure_ctxs[index].psp_start = 0; 
                k_secure_ctxs[index].task_handler = task_handler;
                *(cpu_data_t *)(k_secure_ctxs[index].curr_psp) = K_SECURE_CTX_STACK_SEAL_VALUE;
                *((cpu_data_t *)(k_secure_ctxs[index].curr_psp) + 1) = K_SECURE_CTX_STACK_SEAL_VALUE;
                index += 1;
            } else {
                index = K_SECURE_CTX_INVALID_ID;
            }
        } else {
            index = K_SECURE_CTX_INVALID_ID;
        }
    }
    return index;
}

__NSC__ __KNL__ void tos_secure_ctx_free(k_secure_ctx_handle_t handle, void *task_handler)
{
    uint32_t ulIPSR;
    secure_get_IPSR(ulIPSR);
    if (ulIPSR != 0) {
        if ((handle > 0) && (handle <= TOS_CFG_NUM_SECURE_CTX)) {
            handle -= 1;
            if (k_secure_ctxs[handle].task_handler == task_handler) {
                tos_mmheap_free(k_secure_ctxs[handle].psp_limit);
                k_secure_ctxs[handle].curr_psp = K_NULL;
                k_secure_ctxs[handle].psp_limit = K_NULL;
                k_secure_ctxs[handle].task_handler = K_NULL;
            }
        }
    }
}

__NSC__ __KNL__ void tos_secure_ctx_load(k_secure_ctx_handle_t handle, void *task_handler)
{
    uint32_t ulIPSR;
    uint8_t * ulPSPLimit;
    secure_get_IPSR(ulIPSR);
    if (ulIPSR != 0) {
        if ((handle > 0) && (handle <= TOS_CFG_NUM_SECURE_CTX)) {
            handle -= 1;
            secure_get_PSPLIM(ulPSPLimit);
            if ((ulPSPLimit == K_NULL) && (k_secure_ctxs[handle].task_handler == task_handler)) {
                tos_secure_load_ctx_asm(&(k_secure_ctxs[handle]));
            }
        }
    }
}

__NSC__ __KNL__ void tos_secure_ctx_save(k_secure_ctx_handle_t handle, void *task_handler)
{
    uint32_t ulIPSR;
    uint32_t ulPSPLIM;
    uint32_t ulPSP;
    secure_get_IPSR(ulIPSR);
    if (ulIPSR != 0) {
        if ((handle > 0) && (handle <= TOS_CFG_NUM_SECURE_CTX)) {
            handle -= 1;
            secure_get_PSPLIM(ulPSPLIM);
            if ((ulPSPLIM == (uint32_t)(k_secure_ctxs[handle].psp_limit)) && (k_secure_ctxs[handle].task_handler == task_handler)) {
                tos_secure_save_ctx_asm(&(k_secure_ctxs[handle]));
            }
        }
    }
}
